import tensorflow as tf
import numpy as np
import os as os

from PIL import Image
from src.letters.letter_train import CNN
from src.utils.visualize import *

'''
python 3.7
tensorflow 2.0.0b0
pillow(PIL) 4.3.0
'''


class Predict(object):

    def __init__(self):
        latest = tf.train.latest_checkpoint('./ckpt')
        self.cnn = CNN()
        # 恢复网络权重
        self.cnn.model.load_weights(latest)

    def predict(self, image_path):
        # 读取图片并转换为灰度模式，即将RGB通道合为一个通道，对这个单通道的数据进行操作
        img = Image.open(image_path).convert('L')
        # 预处理，统一转换为28*28大小
        img = prepare_image(img)
        # reshape()改变数组的结构，并且原始数据不发生变化
        # 将这个灰度图片作为一个数组，转换为28大组、28小组、每小组元素个数为1的数组
        # / 255 表示每个数都除以255，将每个数都压缩到0-1之间
        img = np.reshape(img, (28, 28, 1)) / 255
        # 数组的每个数被1减去，并且数组升一维，是为了符合模型的维度
        x = np.array([1 - img])

        # 开始根据模型预测
        y = self.cnn.model.predict(x)

        # 因为x只传入了一张图片，取y[0]即可
        # np.argmax()取得最大值的下标，即代表的数字
        print(image_path)
        print(np.argmax(y[0]))
        print('\t-> 这个图片写的是：', get_mapping(np.argmax(y[0]), with_type="letters"))
        # 可视化的方式展示图片
        show_pred_result(image_path, get_mapping(np.argmax(y[0]), with_type="letters"))


def prepare_image(img: Image) -> Image:
    """
    将待预测的图片进行预处理：
        水平翻转、逆时针旋转90度并统一缩放到28*28px。

    这是因为，EMNIST中的数据集中参与训练的图片，
    是经过水平反转并顺时针旋转90度的，
    所以在预测的时候也要对应地变换一下。
    （所以EMNIST为什么要这么处理图片呢？）
    :param img:
    :return: Image
    """
    return img \
        .transpose(Image.FLIP_LEFT_RIGHT) \
        .transpose(Image.ROTATE_90) \
        .resize((28, 28), Image.ANTIALIAS)


if __name__ == "__main__":
    """
    开始预测
    """
    image_dir = '../../english_images/'
    # cwd = os.getcwd()
    files = os.listdir(image_dir)
    # 遍历批量导入预测文件夹内的所有图片
    app = Predict()
    for image in files:
        app.predict(image_dir + image)
